"""
BEL 引擎核心实现

包含二元错误定位器，严格按照设计文档要求
复用现有的 vLLM 推理框架和统一的工具函数
"""

import json
import re
import os
import sys
from typing import Dict, List, Optional, TypedDict, Any
from pathlib import Path
import time

# 复用现有的推理模块
from vllm import LLM, SamplingParams
from inference.vllm_infer import extract_answer_try_all_methods, template, prompt
from PIL import Image
import jsonlines
from loguru import logger

from .utils import (
    split_steps, insert_labels, create_llm_instance, create_sampling_params,
    validate_answer, normalize_image_path
)


class BinaryErrorLocator:
    """
    二元错误定位器
    
    实现论文 Algorithm 1 的二分查找算法，定位推理路径中的第一个错误步骤
    严格按照设计文档要求，复用统一的工具函数
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化错误定位器
        
        Args:
            config: 配置字典
        """
        self.config = config
        self.llm = None
        
        # 从配置中提取参数
        self.n_mid = config['algorithm']['n_mid']
        self.success_threshold = config['algorithm']['success_threshold']
        
        logger.info(f"BinaryErrorLocator 初始化完成")
    
    def _ensure_llm(self):
        """延迟初始化 LLM 实例，复用统一的初始化函数"""
        if self.llm is None:
            self.llm = create_llm_instance(self.config)
    
    def _test_prefix_correctness(self, question: str, prefix_steps: List[str], 
                               gold_answer: str, image_path: Optional[str] = None) -> float:
        """
        测试前缀的正确性
        
        Args:
            question: 问题文本
            prefix_steps: 前缀步骤列表
            gold_answer: 标准答案
            image_path: 图片路径
            
        Returns:
            float: 正确率 (0.0 ~ 1.0)
        """
        try:
            self._ensure_llm()
            
            # 构建前缀，分为用户输入部分和助手回复部分
            user_prompt = prompt + question
            
            # 构建助手已有的回复部分（前缀步骤）
            assistant_prefix = ""
            if prefix_steps:
                step_parts = []
                for i, step in enumerate(prefix_steps):
                    step_parts.append(f"Step {i+1}: {step}")
                assistant_prefix = '\n'.join(step_parts)
            
            # 使用修改后的 template 格式，让模型从前缀处继续生成
            if assistant_prefix:
                """
                有前缀的情况：在 assistant 部分包含前缀，让模型继续
                格式：<|im_start|>assistant\n前缀步骤内容
                这样模型会从前缀的最后一步继续生成
                <|im_start|>system
                You are a helpful assistant.<|im_end|>
                <|im_start|>user
                <|image|>you are given a math problem image, please solve the problem step by step...
                Question:What is 2+3?<|im_end|>
                <|im_start|>assistant
                Step 1: I need to add 2 and 3
                Step 2: 2 + 3 = 5
                Step 3:
                """
                next_step_num = len(prefix_steps) + 1
                instruction = template.format(user_prompt) + assistant_prefix + f"\nStep {next_step_num}:"
            else:
                # 没有前缀的情况：正常使用 template
                instruction = template.format(user_prompt)
            
            # 处理多模态输入
            inputs = []
            if image_path:
                if os.path.exists(image_path):
                    try:
                        # 尝试打开和处理图片
                        image = Image.open(image_path)
                        if image.mode != 'RGB':
                            image = image.convert('RGB')
                        
                        inputs = [{
                            'prompt': instruction,
                            'multi_modal_data': {
                                'image': image
                            }
                        }]
                        logger.debug(f"成功加载图片: {image_path}")
                    except Exception as e:
                        logger.warning(f"图片处理失败 ({image_path}): {e}")
                        inputs = [instruction]
                else:
                    logger.warning(f"图片文件不存在: {image_path}")
                    inputs = [instruction]
            else:
                inputs = [instruction]
            
            # 创建采样参数，使用正确的方式设置 n_mid
            sampling_params = create_sampling_params(self.config, n=self.n_mid)
            
            # 生成
            try:
                outputs = self.llm.generate(inputs, sampling_params)
            except Exception as e:
                logger.error(f"生成推理路径失败: {e}")
                return 0.0
            
            if not outputs or not outputs[0].outputs:
                logger.warning("未生成任何推理路径")
                return 0.0
            
            # 提取生成的文本并验证
            rollouts = []
            for output in outputs[0].outputs:
                if output.text and output.text.strip():
                    rollouts.append(output.text.strip())
                else:
                    logger.warning("生成了空的推理路径")
            
            if not rollouts:
                logger.warning("所有生成的推理路径都为空")
                return 0.0
            
            # 验证答案，使用统一的答案验证函数
            correct_count = 0
            for rollout in rollouts:
                if validate_answer(rollout, gold_answer):
                    correct_count += 1
            
            # 计算正确率
            success_rate = correct_count / len(rollouts) if rollouts else 0.0
            
            logger.debug(f"前缀测试: {correct_count}/{len(rollouts)} = {success_rate:.2f}")
            return success_rate
            
        except Exception as e:
            logger.error(f"前缀测试失败: {e}")
            return 0.0
    
    def locate_error(self, record: Dict) -> List[List[int]]:
        """
        定位推理路径中的第一个错误步骤
        
        
        Args:
            record: 包含 rollouts 的记录
            
        Returns:
            List[List[int]]: 每个 rollout 的标签列表 (1=正确, 0=错误)
        """
        question = record.get('question', '')
        gold_answer = record.get('gold_answer', '')
        
        relative_image_path = normalize_image_path(record)
        if relative_image_path:
            # 从配置中获取 image_root 并拼接完整路径
            image_root = self.config.get('data', {}).get('image_root', '')
            image_path = os.path.join(image_root, relative_image_path) if image_root else relative_image_path
        else:
            image_path = None
            
        rollouts = record.get('rollouts', [])
        
        if not rollouts:
            return []
        
        all_labels = []
        
        for rollout in rollouts:
            output = rollout.get('output', '')
            steps = split_steps(output)
            
            if not steps:
                all_labels.append([])
                continue
            
            # 初始化标签（全部为1，表示正确）
            labels = [1] * len(steps)
            
            # 首先检查整个推理是否正确
            if validate_answer(output, gold_answer):
                # 如果最终答案正确，所有步骤都标记为正确
                all_labels.append(labels)
                continue
            
            # 二分查找第一个错误步骤
            first_error_idx = self._binary_search_first_error(
                question, steps, gold_answer, image_path
            )
            
            # 标记错误步骤：从第一个错误步骤开始，后续步骤都标记为错误
            for i in range(first_error_idx, len(steps)):
                labels[i] = 0
            
            all_labels.append(labels)
        
        return all_labels
    
    def _binary_search_first_error(self, question: str, steps: List[str], 
                                 gold_answer: str, image_path: Optional[str]) -> int:
        """
        使用二分查找定位第一个错误步骤
        
        Args:
            question: 问题文本
            steps: 推理步骤列表
            gold_answer: 标准答案
            image_path: 图片路径
            
        Returns:
            int: 第一个错误步骤的索引
        """
        if not steps:
            return 0
        
        left, right = 0, len(steps) - 1
        first_error_idx = len(steps)  # 默认为最后一个位置（表示没有找到错误）
        
        # 特殊情况：检查第一个步骤
        if not self._is_prefix_correct(question, steps[:1], gold_answer, image_path):
            return 0
        
        # 特殊情况：检查最后一个步骤
        if self._is_prefix_correct(question, steps, gold_answer, image_path):
            # 如果整个推理都正确，但最终答案错误，可能是答案提取问题
            return len(steps) - 1
        
        # 二分查找
        while left <= right:
            mid = (left + right) // 2
            
            # 测试前缀 [0, mid] 的正确性
            if self._is_prefix_correct(question, steps[:mid + 1], gold_answer, image_path):
                # 前缀正确，错误在后半部分
                left = mid + 1
            else:
                # 前缀错误，记录当前位置并在前半部分继续查找
                first_error_idx = mid
                right = mid - 1
        
        return first_error_idx
    
    def _is_prefix_correct(self, question: str, prefix_steps: List[str], 
                          gold_answer: str, image_path: Optional[str]) -> bool:
        """
        判断前缀是否正确（带缓存优化）
        
        Args:
            question: 问题文本
            prefix_steps: 前缀步骤列表
            gold_answer: 标准答案
            image_path: 图片路径
            
        Returns:
            bool: 前缀是否正确
        """
        success_rate = self._test_prefix_correctness(
            question, prefix_steps, gold_answer, image_path
        )
        return success_rate >= self.success_threshold


def create_default_config(
    model_path: str,
    temperature: float = 0.2,
    top_p: float = 0.9,
    num_rollouts: int = 16,
    max_tokens: int = 2048,
    tensor_parallel_size: int = 1,
    success_threshold: float = 0.6,
    enable_prefix_cache: bool = True
) -> Dict[str, Any]:
    """
    创建默认配置（简化版，移除自适应采样参数）
    
    Args:
        model_path: 模型路径
        temperature: 采样温度
        top_p: 采样参数
        num_rollouts: 推理路径数量
        max_tokens: 最大token数
        tensor_parallel_size: 张量并行大小
        success_threshold: 成功阈值
        enable_prefix_cache: 是否启用前缀缓存
        
    Returns:
        Dict[str, Any]: 配置字典
    """
    return {
        'model': {
            'path': model_path,
            'temperature': temperature,
            'top_p': top_p,
            'max_tokens': max_tokens,
            'tensor_parallel_size': tensor_parallel_size,
        },
        'algorithm': {
            'batch_size': 8,
            'num_rollouts': num_rollouts,
            'n_mid': min(8, num_rollouts),
            'success_threshold': success_threshold,
            'num_samples_per_run': -1,
        },
        'performance': {
            'enable_prefix_caching': enable_prefix_cache,
            'gpu_memory_utilization': 0.85,
            'max_model_len': 4096,
        },
        'output': {
            'rollouts': './outputs/rollouts.jsonl',
            'labeled': './outputs/S_BEL.jsonl',
        }
    } 